package com.example.actordemo3.sharding;

import akka.actor.*;
import akka.cluster.sharding.ClusterSharding;
import akka.cluster.sharding.ClusterShardingSettings;
import akka.cluster.sharding.ShardRegion;
import akka.japi.pf.DeciderBuilder;
import akka.persistence.AbstractPersistentActor;
import scala.PartialFunction;
import scala.Tuple2;
import scala.concurrent.duration.Duration;

import java.util.Optional;


public class ShardingTest {
    ActorSystem system = null;

    ActorRef getSelf() {
        return null;
    }
    static class Counter extends AbstractPersistentActor {

        public enum CounterOp {
            INCREMENT,
            DECREMENT
        }

        public static class Get {
            public final long counterId;

            public Get(long counterId) {
                this.counterId = counterId;
            }
        }

        public static class EntityEnvelope {
            public final long id;
            public final Object payload;

            public EntityEnvelope(long id, Object payload) {
                this.id = id;
                this.payload = payload;
            }
        }

        public static class CounterChanged {
            public final int delta;

            public CounterChanged(int delta) {
                this.delta = delta;
            }
        }

        int count = 0;

        // getSelf().path().name() is the entity identifier (utf-8 URL-encoded)
        @Override
        public String persistenceId() {
            return "Counter-" + getSelf().path().name();
        }

        @Override
        public void preStart() throws Exception {
            super.preStart();
            getContext().setReceiveTimeout(Duration.create(120,"s"));
        }

        void updateState(CounterChanged event) {
            count += event.delta;
        }

        @Override
        public Receive createReceiveRecover() {
            return receiveBuilder().match(CounterChanged.class, this::updateState).build();
        }

        @Override
        public Receive createReceive() {
            return receiveBuilder()
                    .match(Get.class, this::receiveGet)
                    .matchEquals(CounterOp.INCREMENT, msg -> receiveIncrement())
                    .matchEquals(CounterOp.DECREMENT, msg -> receiveDecrement())
                    .matchEquals(ReceiveTimeout.getInstance(), msg -> passivate())
                    .build();
        }

        private void receiveGet(Get msg) {
            getSender().tell(count, getSelf());
        }

        private void receiveIncrement() {
            persist(new CounterChanged(+1), this::updateState);
        }

        private void receiveDecrement() {
            persist(new CounterChanged(-1), this::updateState);
        }

        private void passivate() {
            getContext().getParent().tell(new ShardRegion.Passivate(PoisonPill.getInstance()), getSelf());
        }
    }


    public  void demonstrateUsage(){
        ShardRegion.MessageExtractor messageExtractor = new ShardRegion.MessageExtractor() {
            @Override
            public String shardId(Object message) {
                int numberOfShards = 100;
                if (message instanceof Counter.EntityEnvelope) {
                    long id = ((Counter.EntityEnvelope) message).id;
                    return String.valueOf(id % numberOfShards);
                } else if (message instanceof Counter.Get) {
                    long id = ((Counter.Get) message).counterId;
                    return String.valueOf(id % numberOfShards);
                } else {
                    return null;
                }
            }

            @Override
            public String entityId(Object message) {
                if (message instanceof Counter.EntityEnvelope)
                    return String.valueOf(((Counter.EntityEnvelope) message).id);
                else if (message instanceof Counter.Get)
                    return String.valueOf(((Counter.Get) message).counterId);
                else return null;
            }

            @Override
            public Object entityMessage(Object message) {
                if (message instanceof Counter.EntityEnvelope)
                    return ((Counter.EntityEnvelope) message).payload;
                else return message;
            }
        };

        ClusterShardingSettings settings = ClusterShardingSettings.create(system);
        ActorRef startedCounterRegion =
                ClusterSharding.get(system)
                        .start("Counter", Props.create(Counter.class), settings,messageExtractor );
        ActorRef counterRegion = ClusterSharding.get(system).shardRegion("Counter");
        counterRegion.tell(new Counter.Get(123), getSelf());

        counterRegion.tell(new Counter.EntityEnvelope(123, Counter.CounterOp.INCREMENT), getSelf());
        counterRegion.tell(new Counter.Get(123), getSelf());
        // #counter-usage

        // #counter-supervisor-start
        ClusterSharding.get(system)
                .start(
                        "SupervisedCounter", Props.create(CounterSupervisor.class), settings, messageExtractor);
        // #counter-supervisor-start

//        // #proxy-dc
//        ActorRef counterProxyDcB =
//                ClusterSharding.get(system)
//                        .startProxy(
//                                "Counter",
//                                Optional.empty().get(),
//                                new PartialFunction<Object, Tuple2<String, Object>>() {
//                                }
//                                Optional.of("B"), // data center name
//                                messageExtractor);
    }

    public void demonstrateUsage2() {
        ShardRegion.MessageExtractor messageExtractor =
                new ShardRegion.MessageExtractor() {

                    @Override
                    public String entityId(Object message) {
                        if (message instanceof Counter.EntityEnvelope)
                            return String.valueOf(((Counter.EntityEnvelope) message).id);
                        else if (message instanceof Counter.Get)
                            return String.valueOf(((Counter.Get) message).counterId);
                        else return null;
                    }

                    @Override
                    public Object entityMessage(Object message) {
                        if (message instanceof Counter.EntityEnvelope)
                            return ((Counter.EntityEnvelope) message).payload;
                        else return message;
                    }

                    // #extractShardId-StartEntity
                    @Override
                    public String shardId(Object message) {
                        int numberOfShards = 100;
                        if (message instanceof Counter.EntityEnvelope) {
                            long id = ((Counter.EntityEnvelope) message).id;
                            return String.valueOf(id % numberOfShards);
                        } else if (message instanceof Counter.Get) {
                            long id = ((Counter.Get) message).counterId;
                            return String.valueOf(id % numberOfShards);
                        } else if (message instanceof ShardRegion.StartEntity) {
                            long id = Long.valueOf(((ShardRegion.StartEntity) message).entityId());
                            return String.valueOf(id % numberOfShards);
                        } else {
                            return null;
                        }
                    }
                    // #extractShardId-StartEntity

                };
    }

    // #supervisor
    public  static class CounterSupervisor extends AbstractActor {

        private final ActorRef counter =
                getContext().actorOf(Props.create(Counter.class), "theCounter");

        private static final SupervisorStrategy strategy =
                new OneForOneStrategy(
                        DeciderBuilder.match(IllegalArgumentException.class, e -> SupervisorStrategy.resume())
                                .match(ActorInitializationException.class, e -> SupervisorStrategy.stop())
                                .match(Exception.class, e -> SupervisorStrategy.restart())
                                .matchAny(o -> SupervisorStrategy.escalate())
                                .build());

        @Override
        public SupervisorStrategy supervisorStrategy() {
            return strategy;
        }

        @Override
        public Receive createReceive() {
            return receiveBuilder()
                    .match(Object.class, msg -> counter.forward(msg, getContext()))
                    .build();
        }
    }

}
